#include <unistd.h>
#include <inttypes.h>
#include <fstream>
#include "callgraph.h"

namespace hfsort
{

  constexpr uint32_t BUFLEN = 24000;
  constexpr uint32_t kPageSize = 2 << 20;

  void error(const char *msg)
  {
    printf("ERROR: %s\n", msg);
    exit(1);
  }

  void readSymbols(MyCallGraph &cg, std::string binary)
  {
    uint64 addr = InvalidAddr;
    uint64 size = 0;
    std::string name = "";
    AddressSymbolMap address_symbol_map_;

    ElfReader elf_reader(binary);
    // uint64 base_addr_ = elf_reader.VaddrOfFirstLoadSegment();

    SymbolReader symbol_reader(&address_symbol_map_);
    elf_reader.VisitSymbols(&symbol_reader);
    
    for (const auto &addr_symbol : address_symbol_map_) {
      addr = addr_symbol.first;
      name = addr_symbol.second.first;
      size = addr_symbol.second.second;
      cg.addFunc(name, addr, size, 0);
      // majin_info("name: %s, addr: %lx, size: %x",name.c_str(), addr, size);
    }
  }

  uint64_t getAddr(char *perfLine)
  {
    uint64_t addr;
    int ret = sscanf(perfLine, "%" SCNx64, &addr);
    if (ret != 1)
      return InvalidAddr;
    return addr;
  }

  TargetId getTargetId(const MyCallGraph &cg, uint64_t addr)
  {
    if (addr == InvalidAddr)
    {
      majin_debug("get InvalidAddr");
      return InvalidId;
    }
    return cg.addrToTargetId(addr);
  }

  void readPerfData(MyCallGraph &cg, FILE *file, bool computeArcWeight)
  {
    char line[BUFLEN];

    while (fgets(line, BUFLEN, file) != NULL)
    {
      if (line[0] == '#')
        continue;
      if (isspace(line[0]))
        continue;

      // process one sample
      if (fgets(line, BUFLEN, file) == NULL)
        error("reading perf data");
      auto addrTop = getAddr(line);
      TargetId idTop = getTargetId(cg, addrTop);
      if (idTop == InvalidId)
        continue;
      cg.targets[idTop].samples++;
      majin_debug("readPerfData: idTop: %u %s\n", idTop,
                  cg.funcs[idTop].mangledNames[0].c_str());
      if (fgets(line, BUFLEN, file) == NULL)
        error("reading perf data");
      auto addrCaller = getAddr(line);
      TargetId idCaller = getTargetId(cg, addrCaller);
      if (idCaller != InvalidId)
      {
        auto &arc = cg.incArcWeight(idCaller, idTop, computeArcWeight ? 1 : 0);
        if (computeArcWeight)
        {
          arc.avgCallOffset += addrCaller - cg.funcs[idCaller].addr;
        }
        majin_debug("readPerfData: idCaller: %u %s\n", idCaller,
                    cg.funcs[idCaller].mangledNames[0].c_str());
      }
    }

    if (!computeArcWeight)
      return;

    cg.normalizeArcWeights();
  }

  std::string getNameWithoutSuffix(std::string str)
  {
    int suffixStartPosition = str.find(".");
    if (suffixStartPosition == -1)
    {
      // if no suffix is found, just add the wildcard
      return str + "*";
    }
    else
    {
      // replace sufix with wildcard
      return str.substr(0, suffixStartPosition) + "*";
    }
  }

  void print(MyCallGraph &cg, const char *filename,
             const std::vector<Cluster> &clusters, bool useWildcards)
  {
    FILE *outfile = fopen(filename, "wt");
    if (!outfile)
    {
      error("opening output file faild !");
    }
    uint32_t totalSize = 0;
    uint32_t curPage = 0;
    uint32_t hotfuncs = 0;
    double totalDistance = 0;
    double totalCalls = 0;
    double totalCalls64B = 0;
    double totalCalls4KB = 0;
    double totalCalls32KB = 0;
    double totalCalls128KB = 0;
    double totalCalls512KB = 0;
    double totalCalls2MB = 0;
    std::unordered_map<TargetId, uint64_t> newAddr;
    for (auto &cluster : clusters)
    {
      for (auto fid : cluster.targets)
      {
        if (cg.targets[fid].samples > 0)
        {
          newAddr[fid] = totalSize;
          totalSize += cg.targets[fid].size;
        }
      }
    }
    totalSize = 0;
    majin_debug("============== page 0 ==============\n");
    for (auto &cluster : clusters)
    {
      majin_debug(
          "-------- density = %.3lf (%u / %u) --------\n",
          (double)cluster.samples / cluster.size,
          cluster.samples, cluster.size);
      for (auto fid : cluster.targets)
      {
        if (cg.targets[fid].samples > 0)
        {
          hotfuncs++;
          int space = 0;
          for (const auto &mangledName : cg.funcs[fid].mangledNames)
          {
            if (useWildcards)
            {
              fprintf(outfile, "%.*s.text.%s\n",
                      space, " ", getNameWithoutSuffix(mangledName).c_str());
            }
            else
            {
              fprintf(outfile, "%.*s.text.%s\n",
                      space, " ", mangledName.c_str());
            }

            space = 1;
          }
          uint64_t dist = 0;
          uint64_t calls = 0;
          for (auto dst : cg.targets[fid].succs)
          {
            auto &arc = *cg.arcs.find(Arc(fid, dst));
            auto d = std::abs(newAddr[arc.dst] -
                              (newAddr[fid] + arc.avgCallOffset));
            auto w = arc.weight;
            calls += w;
            if (d < 64)
              totalCalls64B += w;
            if (d < 4096)
              totalCalls4KB += w;
            if (d < 1 << 15)
              totalCalls32KB += w;
            if (d < 1 << 17)
              totalCalls128KB += w;
            if (d < 1 << 19)
              totalCalls512KB += w;
            if (d < 1 << 21)
              totalCalls2MB += w;
            majin_debug(
                "arc: %u [@%lu+%.1lf] -> %u [@%lu]: weight = %.0lf, "
                "callDist = %f\n",
                arc.src, newAddr[arc.src], arc.avgCallOffset,
                arc.dst, newAddr[arc.dst], arc.weight, d);
            dist += arc.weight * d;
          }
          totalCalls += calls;
          totalDistance += dist;
          majin_debug("start = %6u : avgCallDist = %lu : %s\n",
                      totalSize,
                      calls ? dist / calls : 0,
                      cg.toString(fid).c_str());
          totalSize += cg.targets[fid].size;
          uint32_t newPage = totalSize / kPageSize;
          if (newPage != curPage)
          {
            curPage = newPage;
            majin_debug("============== page %u ==============\n", curPage);
          }
        }
      }
    }
    fclose(outfile);
    printf("Output saved in file %s\n", filename);
    printf("  Number of hot functions: %u\n  Number of clusters: %lu\n",
           hotfuncs, clusters.size());
    printf("  Final average call distance = %.1lf (%.0lf / %.0lf)\n",
           totalCalls ? totalDistance / totalCalls : 0,
           totalDistance, totalCalls);
    printf("  Total Calls = %.0lf\n", totalCalls);
    if (totalCalls)
    {
      printf("  Total Calls within 64B = %.0lf (%.2lf%%)\n",
             totalCalls64B, 100 * totalCalls64B / totalCalls);
      printf("  Total Calls within 4KB = %.0lf (%.2lf%%)\n",
             totalCalls4KB, 100 * totalCalls4KB / totalCalls);
      printf("  Total Calls within 32KB = %.0lf (%.2lf%%)\n",
             totalCalls32KB, 100 * totalCalls32KB / totalCalls);
      printf("  Total Calls within 128KB = %.0lf (%.2lf%%)\n",
             totalCalls128KB, 100 * totalCalls128KB / totalCalls);
      printf("  Total Calls within 512KB = %.0lf (%.2lf%%)\n",
             totalCalls512KB, 100 * totalCalls512KB / totalCalls);
      printf("  Total Calls within 2MB = %.0lf (%.2lf%%)\n",
             totalCalls2MB, 100 * totalCalls2MB / totalCalls);
    }
  }

}

bool MatchBinary(const string &name, const string &full_name){
  return full_name == basename(name.c_str());
}

int main(int argc, char *argv[])
{

  using namespace hfsort;

  MyCallGraph cg;

  const char *symbFileName = NULL;
  const char *perfFileName = NULL;

  FILE *perfFile = NULL;
  bool useWildcards = false;

  Algorithm algorithm = Algorithm::Hfsort;

  extern char *optarg;
  extern int optind;
  int c;
  while ((c = getopt(argc, argv, "pae:s:w")) != -1)
  {
    switch (c)
    {
    case 'p':
      algorithm = Algorithm::PettisHansen;
      break;
    case 'a':
      algorithm = Algorithm::HfsortPlus;
      break;
    case 'w':
      useWildcards = true;
      break;
    case '?':
      error("Unsupported command line argument");
    }
  }

  if (argc < optind + 1 || argc > optind + 2)
  {
    error(
        "Usage: hfsort [-p] [-e <EDGCNT_FILE>] [-s <SYMBOLCNT_FILE>] [-w]"
        " <SYMBOL_FILE> [<PERF_DATA_FILE>]\n"
        "   -p,               use pettis-hansen algorithm for code layout\n"
        "   -a,               use hfsort-plus algorithm for code layout\n"
        "   -w                use wildcards instead of suffixes in function names");
  }

  symbFileName = argv[optind];
  if (argc == optind + 2)
  {
    perfFileName = argv[optind + 1];
  }

  if (symbFileName == NULL)
  {
    error("Error opening symbol file\n");
  }
  readSymbols(cg, symbFileName);

  std::string perCmd = string("perf script -i ") + perfFileName  + 
        string(" --fields comm,pid,ip | sed -ne \"/^\\w*") +
        basename(symbFileName) + string(".*\\$/,+2p\"");

  majin_info("perCmd: %s", perCmd.c_str());

  if (perfFileName && !(perfFile = popen(perCmd.c_str(), "r"))){
    error("Error parser perf data file\n");
  }

  readPerfData(cg, perfFile, true);

  std::vector<Cluster> clusters;
  const char *filename;
  if (algorithm == Algorithm::Hfsort)
  {
    majin_info("=== algorithm : hfsort");
    clusters = clusterize(cg);
    filename = "hotfuncs.txt";
  }
  else if (algorithm == Algorithm::HfsortPlus)
  {
    majin_info("=== algorithm : hfsort-plus");
    clusters = hfsortPlus(cg);
    filename = "hotfuncs-plus.txt";
  }
  else if (algorithm == Algorithm::PettisHansen)
  {
    majin_info("=== algorithm : pettis-hansen");
    clusters = pettisAndHansen(cg);
    filename = "hotfuncs-ph.txt";
  }
  else
  {
    error("Unknown layout algorithm\n");
  }

  sort(clusters.begin(), clusters.end(), compareClustersDensity);
  print(cg, filename, clusters, useWildcards);

  return 0;
}